import json
import cv2
import numpy as np
import matplotlib.pyplot as plt
import dlib
 

def get_point(key_points,index):
    return (int(key_points[index*3]),int(key_points[index*3+1]))

def change_color(path,outPath,name):
    img0 = cv2.imread(path+name+".jpg")
    img1 = cv2.resize(img0, dsize = None, fx = 0.5, fy = 0.5)
    img2 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)

    with open(path+name+"_keypoints.json") as f:
        line=f.readline()
        json_obj=json.loads(line)
        keypoints=json_obj["people"][0]["pose_keypoints_2d"]
    point=get_point(keypoints,0)
    point=(int(point[0]/2),int(point[1]/2))
    p1=get_point(keypoints,1)
    p8=get_point(keypoints,8)
    val=img2[point[1],point[0]]
    diff=6
    # dfs
    points=[]
    queue=[]
    queue.append(((p1[0]+p8[0])//4,(p1[1]+p8[1])//4))
    count=0
    while(len(queue)!=0):
        p=queue[0]
        points.append(p)
        queue.remove(p)
        val=img2[p[1],p[0]]
        print(p)
        if(p[0]<0 or p[1]<0 or p[0]>=img1.shape[1] or p[1]>=img1.shape[0]):
            continue
        if(not (p[0],p[1]+1) in points and not (p[0],p[1]+1) in queue and p[1]<img1.shape[0]-1 and img2[p[1]+1,p[0]]>val-diff and img2[p[1]+1,p[0]]<val+diff):
            queue.append((p[0],p[1]+1))
            points.append((p[0],p[1]+1))
        if(not (p[0],p[1]-1) in points and not (p[0],p[1]-1) in queue and p[1]>0 and img2[p[1]-1,p[0]]>val-diff and img2[p[1]-1,p[0]]<val+diff):
            queue.append((p[0],p[1]-1))
            points.append((p[0],p[1]-1))
        if(not (p[0]+1,p[1]) in points and not (p[0]+1,p[1]) in queue and  p[0]<img1.shape[1]-1 and img2[p[1],p[0]+1]>val-diff and img2[p[1],p[0]+1]<val+diff):
            queue.append((p[0]+1,p[1]))
            points.append((p[0]+1,p[1]))
        if(not (p[0]-1,p[1]) in points and not (p[0]-1,p[1]) in queue and p[0]>0 and img2[p[1],p[0]-1]>val-diff and img2[p[1],p[0]-1]<val+diff):
            queue.append((p[0]-1,p[1]))
            points.append((p[0]-1,p[1]))
        count+=1


# cv2.circle(img=img2, center=(293, 295), radius=5, color=(0, 255, 0), thickness=-1)
    h, w = img1.shape[:2]
    print(h, w)
    for point in points:
        img1[point[1],point[0],0]=0xff
        img1[point[1],point[0],1]=0xcc
        img1[point[1],point[0],2]=0x66
    cv2.imwrite(outPath+name+".jpg",img1)
    cv2.waitKey(delay = 0)

if __name__=="__main__":
    path="./graph/"
    outPath="./change_color/"
    for i in range(1,6):
        change_color(path,outPath,"0"+str(i))